# -*- coding: utf-8 -*-
# import the necessary packages
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F
import numpy as np
import gym
from matplotlib import pyplot as plt
import matplotlib
import math
import os.path


BATCH_SIZE = 32     # batch size of sampling process from buffer
LR = 0.01           # learning rate
EPSILON = 0.9       # epsilon used for epsilon greedy approach
GAMMA = 0.9         # discount factor
TARGET_NETWORK_REPLACE_FREQ = 500       # How frequently target netowrk updates
MEMORY_CAPACITY = 2000                  # The capacity of experience replay buffer
DATA_DIR = 'reinforcement-learning/DQN/data/'

env = gym.make("CartPole-v0")
env = env.unwrapped
N_ACTIONS = env.action_space.n  # 2 actions
N_STATES = env.observation_space.shape[0] # 4 states
ENV_A_SHAPE = 0 if isinstance(env.action_space.sample(), int) else env.action_space.sample().shape    

class Net(nn.Module):
    def __init__(self):
        # Define the network structure, a very simple fully connected network
        super(Net, self).__init__()
        # Define the structure of fully connected network
        self.fc1 = nn.Linear(N_STATES, 10)  # layer 1
        self.fc1.weight.data.normal_(0, 0.1) # in-place initilization of weights of fc1
        self.out = nn.Linear(10, N_ACTIONS) # layer 2
        self.out.weight.data.normal_(0, 0.1) # in-place initilization of weights of fc2
        
        
    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        actions_value = self.out(x)
        return actions_value

class DQN(object):
    def __init__(self):
        self.load_model()
        self.learn_step_counter = 0
        self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=LR)
        self.loss_func = nn.MSELoss()

    def choose_action(self, x):
        x = torch.unsqueeze(torch.FloatTensor(x), 0)
        if self.memory_counter >= MEMORY_CAPACITY and np.random.uniform() < EPSILON:
            action_value = self.eval_net.forward(x)
            action = torch.max(action_value, 1)[1].data.numpy()
            action = action[0] if ENV_A_SHAPE == 0 else action.reshape(ENV_A_SHAPE) 
        else:
            action = np.random.randint(0, N_ACTIONS)
            action = action if ENV_A_SHAPE == 0 else action.reshape(ENV_A_SHAPE)

        return action
    

    def store_transition(self, s, a, r, s_):
        transition = np.hstack((s, [a, r], s_))
        index = self.memory_counter % MEMORY_CAPACITY
        self.memory[index, :] = transition
        self.memory_counter += 1


    def learn(self, remeber = False):
        if remeber or self.learn_step_counter % TARGET_NETWORK_REPLACE_FREQ == 0:
            self.target_net.load_state_dict(self.eval_net.state_dict())
        self.learn_step_counter += 1
        
        sample_index = np.random.choice(MEMORY_CAPACITY, BATCH_SIZE)
        b_memory = self.memory[sample_index, :]
        b_s = Variable(torch.FloatTensor(b_memory[:, :N_STATES]))
        b_a = Variable(torch.LongTensor(b_memory[:, N_STATES:N_STATES+1].astype(int)))
        b_r = Variable(torch.FloatTensor(b_memory[:, N_STATES+1:N_STATES+2]))
        b_s_ = Variable(torch.FloatTensor(b_memory[:, -N_STATES:]))
        
        q_eval = self.eval_net(b_s).gather(1, b_a) 
        q_next = self.target_net(b_s_).detach()
        
        q_target = b_r + GAMMA * q_next.max(1)[0].view(BATCH_SIZE, 1)
        loss = self.loss_func(q_eval, q_target)
        
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

    def save_model(self):
        torch.save(self.eval_net.state_dict(), DATA_DIR + 'eval.pkl')
        torch.save(self.target_net.state_dict(), DATA_DIR + 'target.pkl')
        np.save(DATA_DIR + 'memory', self.memory)

    def load_model(self):
        self.eval_net = Net()
        self.target_net = Net()
        self.memory = np.zeros((MEMORY_CAPACITY, N_STATES * 2 + 2))
        self.memory_counter = 0
        if os.path.isfile(DATA_DIR + 'eval.pkl'):
            self.eval_net.load_state_dict(torch.load(DATA_DIR + 'eval.pkl'))
        if os.path.isfile(DATA_DIR + 'target.pkl'):
            self.target_net.load_state_dict(torch.load(DATA_DIR + 'target.pkl'))
        if os.path.isfile(DATA_DIR + 'memory.npy'):
            self.memory = np.load(DATA_DIR + 'memory.npy')
            self.memory_counter = len(self.memory)


dqn = DQN()

plt_x = []
plt_y = []

for i_episode in range(1000):
    s = env.reset()
    ep_r = 0
    while True:
        # env.render()
        a = dqn.choose_action(s)
        s_, r, done, info = env.step(a)
        
        x, x_dot, theta, theta_dot = s_
        # r1 = (env.x_threshold - abs(x)) / env.x_threshold - 0.8
        # r2 = (env.theta_threshold_radians - abs(theta)) / env.theta_threshold_radians - 0.5

        # r1代表车的 x水平位移 与 x最大边距 的距离差的得分
        r1 = math.exp((env.x_threshold - abs(x))/env.x_threshold) - math.exp(1)/2
        # r2代表棒子的 theta离垂直的角度 与 theta最大角度 的差的得分
        r2 = math.exp((env.theta_threshold_radians - abs(theta))/env.theta_threshold_radians) - math.exp(1)/2
        r = r1 + r2
        
        dqn.store_transition(s, a, r, s_)
        
        ep_r += r
        if dqn.memory_counter >= MEMORY_CAPACITY and done:
            dqn.learn(False)
            plt_x.append(i_episode)
            plt_y.append(round(ep_r, 2))
            print('Ep: ', i_episode, ' |', 'Ep_r: ', round(ep_r, 2))
            break
        
        if done:
            break
        s = s_ 

dqn.save_model()
plt.scatter(plt_x, plt_y)
plt.show()